{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You may need to restart your runtime prior to this, to let your installation take effect\n",
    "# Some basic setup:\n",
    "# Setup detectron2 logger\n",
    "import detectron2\n",
    "from detectron2.utils.logger import setup_logger\n",
    "setup_logger()\n",
    "\n",
    "# import some common libraries\n",
    "import numpy as np\n",
    "import os\n",
    "import cv2\n",
    "import random\n",
    "# from google.colab.patches import cv2_imshow\n",
    "\n",
    "# import some common detectron2 utilities\n",
    "from detectron2 import model_zoo\n",
    "from detectron2.engine import DefaultPredictor\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "from detectron2.data import DatasetCatalog,MetadataCatalog\n",
    "from detectron2.engine import DefaultTrainer\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.pylab as pylab\n",
    "pylab.rcParams['figure.figsize'] = 10,10\n",
    "def imshow(img):\n",
    "    plt.imshow(img[:, :, [2, 1, 0]])\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train ORCNN on amodal datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register my amodal datasets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from detectron2.data.datasets import register_coco_instances\n",
    "from detectron2.data import MetadataCatalog\n",
    "dataDir='datasets/coco'\n",
    "annTrainFile='{}/annotations/COCO_amodal_train2014_with_classes_poly.json'.format(dataDir)\n",
    "imgTrainFile = '{}/train2014'.format(dataDir)\n",
    "register_coco_instances(\"amodal_coco_train\", {},annTrainFile , imgTrainFile)\n",
    "# Prepare test datasets \n",
    "annTestFile='{}/annotations/COCO_amodal_val2014_with_classes_poly.json'.format(dataDir)\n",
    "imgTestFile = '{}/val2014'.format(dataDir)\n",
    "from detectron2.data.datasets import register_coco_instances\n",
    "from detectron2.data import MetadataCatalog\n",
    "register_coco_instances(\"amodal_coco_val\", {}, annTestFile, imgTestFile)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dicts = DatasetCatalog.get(\"amodal_coco_train\")\n",
    "annos = [img['annotations'] for img in dataset_dicts ]\n",
    "invisible_mask = [ obj['invisible_mask'] for img in annos for obj in img if obj.get('invisible_mask')]\n",
    "visible_mask = [ obj['visible_mask'] for img in annos for obj in img ]\n",
    "print(invisible_mask[2])\n",
    "[len(i) for i in invisible_mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity checks : overfitting small datasets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Register small datasets for debugging \n",
    "dataset_dicts = DatasetCatalog.get(\"amodal_coco_train\")\n",
    "import random\n",
    "imgs = random.sample(dataset_dicts,k=10)\n",
    "# imgs = dataset_dicts[:10]\n",
    "DatasetCatalog.register(\"small_amodal_test\", lambda : imgs)\n",
    "metadata = {}\n",
    "MetadataCatalog.get(\"small_amodal_test\").set(\n",
    "        image_root=\"datasets/coco/train2014\", evaluator_type=\"coco\", **metadata\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycocotools.coco import COCO\n",
    "coco_api = COCO(annTrainFile)\n",
    "cat_ids = sorted(coco_api.getCatIds())\n",
    "cats = coco_api.loadCats(cat_ids)\n",
    "# The categories in a custom json file may not be sorted.\n",
    "thing_classes = [c[\"name\"] for c in sorted(cats, key=lambda x: x[\"id\"])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = {}\n",
    "MetadataCatalog.get(\"small_amodal_test\").set(thing_classes = thing_classes, **metadata )\n",
    "# MetadataCatalog.get(\"small_amodal_test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smalldicts = DatasetCatalog.get(\"small_amodal_test\")\n",
    "for i in range(len(smalldicts)):\n",
    "    im = cv2.imread(smalldicts[i][\"file_name\"])\n",
    "    imshow(im[:, :, ::-1])\n",
    "    visualizer = Visualizer(im[:, :, ::-1], metadata=MetadataCatalog.get(\"small_amodal_test\"), scale=0.5)\n",
    "    vis = visualizer.draw_dataset_dict(smalldicts[i])\n",
    "    imshow(vis.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training \n",
    "- remove orcnnEvaluation directory everytime : create json_file for small dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = get_cfg()\n",
    "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_orcnn_R_50_FPN_3x.yaml\"))\n",
    "cfg.DATASETS.TRAIN = (\"small_amodal_test\",)\n",
    "cfg.DATASETS.TEST = (\"small_amodal_test\",)\n",
    "cfg.DATALOADER.NUM_WORKERS = 2\n",
    "# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\")  # Let training initialize from model zoo\n",
    "cfg.SOLVER.IMS_PER_BATCH = 2\n",
    "cfg.SOLVER.BASE_LR = 0.0015  # pick a good LR\n",
    "# cfg.SOLVER.STEPS = (1300,1400)\n",
    "cfg.SOLVER.MAX_ITER = 1000\n",
    "cfg.VIS_PERIOD = 20\n",
    "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 80\n",
    "cfg.OUTPUT_DIR = \"orcnnCheckpoint\"\n",
    "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9  # set the testing threshold for this model\n",
    "trainer = DefaultTrainer(cfg) \n",
    "trainer.resume_or_load(resume=False)\n",
    "# import pdb; pdb.set_trace()\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tensorboard --logdir=orcnnCheckpoint --port=6006"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9   # set the testing threshold for this model\n",
    "#  evaluate its performance using AP metric implemented in COCO API.\n",
    "from detectron2.evaluation import AmodalEvaluator, inference_on_dataset\n",
    "from detectron2.data import build_detection_test_loader\n",
    "evaluator = AmodalEvaluator(\"small_amodal_test\", cfg, False, output_dir=\"orcnnEvaluation\")\n",
    "val_loader = build_detection_test_loader(cfg, \"small_amodal_test\")\n",
    "# import pdb;pdb.set_trace()\n",
    "inference_on_dataset(trainer.model, val_loader, evaluator) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from detectron2.utils.visualizer import ColorMode\n",
    "import random\n",
    "from detectron2.data import DatasetCatalog\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9   # set the testing threshold for this model\n",
    "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
    "cfg.DATASETS.TEST = (\"small_amodal_test\",)\n",
    "predictor = DefaultPredictor(cfg)\n",
    "dataset_dicts = DatasetCatalog.get(\"small_amodal_test\")\n",
    "# for d in random.sample(dataset_dicts, 2): \n",
    "for i,d in enumerate(dataset_dicts):    \n",
    "    im = cv2.imread(d[\"file_name\"])\n",
    "    visualizer = Visualizer(im[:, :, ::-1], metadata=MetadataCatalog.get(\"small_amodal_test\"), scale=0.5)\n",
    "    vis = visualizer.draw_dataset_dict(d,'invisible_mask')\n",
    "    imshow(vis.get_image()[:, :, ::-1])\n",
    "#     import pdb;pdb.set_trace()\n",
    "    outputs = predictor(im)\n",
    "    v = Visualizer(im[:, :, ::-1],\n",
    "                   metadata=MetadataCatalog.get(\"small_amodal_test\"), \n",
    "                   scale=0.8, \n",
    "                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels\n",
    "    )\n",
    "    v = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"),'pred_invisible_masks')\n",
    "    imshow(v.get_image()[:, :, ::-1])\n",
    "    cv2.imwrite('SmallSamplesVisualizing/Sample{}.png'.format(i), vis.get_image()[:, :, ::-1]) \n",
    "    cv2.imwrite('SmallSamplesVisualizing/Test{}.png'.format(i), v.get_image()[:, :, ::-1]) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training from a COCO-pretrained model as provided by Detectron2\n",
    "1. Finetuning : For the case of COCOA amodal the final output layers that are class-specific had to be initialized randomly as the number of classes and their semantic meaning did not fit to the number of classes of COCO."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from detectron2.engine import DefaultTrainer\n",
    "from detectron2.config import get_cfg\n",
    "\n",
    "cfg = get_cfg()\n",
    "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_orcnn_R_50_FPN_3x.yaml\"))\n",
    "cfg.DATASETS.TRAIN = (\"amodal_coco_train\",)\n",
    "cfg.DATASETS.TEST = (\"amodal_coco_val\",)\n",
    "cfg.DATALOADER.NUM_WORKERS = 2\n",
    "cfg.SOLVER.IMS_PER_BATCH = 2\n",
    "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 80\n",
    "cfg.SOLVER.BASE_LR = 0.0005  # pick a good LR\n",
    "cfg.SOLVER.STEPS = (6000,8000)\n",
    "cfg.SOLVER.MAX_ITER = 10000 \n",
    "cfg.VIS_PERIOD = 500\n",
    "cfg.OUTPUT_DIR = \"orcnnCheckpoint\"\n",
    "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9  # set the testing threshold for this model\n",
    "import pdb; pdb.set_trace()\n",
    "trainer = DefaultTrainer(cfg) \n",
    "trainer.resume_or_load(resume=False)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tensorboard --logdir=orcnnCheckpoint --port=6006"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evalution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg.DATASETS.TEST = (\"amodal_coco_val\",)\n",
    "cfg.DATASETS.TRAIN = (\"amodal_coco_train\",)\n",
    "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9   # set the testing threshold for this model\n",
    "#  evaluate its performance using AP metric implemented in COCO API.\n",
    "from detectron2.evaluation import AmodalEvaluator, inference_on_dataset\n",
    "from detectron2.data import build_detection_test_loader\n",
    "evaluator = AmodalEvaluator(\"amodal_coco_val\", cfg, False, output_dir=\"orcnnEvaluation\")\n",
    "val_loader = build_detection_test_loader(cfg, \"amodal_coco_val\")\n",
    "inference_on_dataset(trainer.model, val_loader, evaluator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from detectron2.utils.visualizer import ColorMode\n",
    "import random\n",
    "from detectron2.data import DatasetCatalog\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "cfg = get_cfg()\n",
    "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_orcnn_R_50_FPN_3x.yaml\"))\n",
    "cfg.OUTPUT_DIR = \"orcnnCheckpoint\"\n",
    "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 80\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9   # set the testing threshold for this model\n",
    "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
    "cfg.DATASETS.TEST = (\"amodal_coco_val\",)\n",
    "predictor = DefaultPredictor(cfg)\n",
    "dataset_dicts = DatasetCatalog.get(\"amodal_coco_val\")\n",
    "for d in random.sample(dataset_dicts, 1):    \n",
    "    im = cv2.imread(d[\"file_name\"])\n",
    "    visualizer = Visualizer(im[:, :, ::-1], metadata=MetadataCatalog.get(\"amodal_coco_val\"), scale=0.5)\n",
    "    vis = visualizer.draw_dataset_dict(d,\"invisible_mask\")\n",
    "    imshow(vis.get_image()[:, :, ::-1])\n",
    "#     import pdb;pdb.set_trace()\n",
    "    outputs = predictor(im)\n",
    "    v = Visualizer(im[:, :, ::-1],\n",
    "                   metadata=MetadataCatalog.get(\"amodal_coco_val\"), \n",
    "                   scale=0.8, \n",
    "                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels\n",
    "    )\n",
    "    v = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"),'pred_invisible_masks')\n",
    "    imshow(v.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = cv2.imread(d[\"file_name\"])\n",
    "visualizer = Visualizer(im[:, :, ::-1], metadata=MetadataCatalog.get(\"amodal_coco_val\"), scale=0.5)\n",
    "vis = visualizer.draw_dataset_dict(d,\"invisible_mask\")\n",
    "imshow(vis.get_image()[:, :, ::-1])\n",
    "#     import pdb;pdb.set_trace()\n",
    "outputs = predictor(im)\n",
    "v = Visualizer(im[:, :, ::-1],\n",
    "           metadata=MetadataCatalog.get(\"amodal_coco_val\"), \n",
    "           scale=0.8, \n",
    "           instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels\n",
    ")\n",
    "v = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"),'pred_invisible_masks')\n",
    "imshow(v.get_image()[:, :, ::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
